# Copyright 2022 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""generate gaussian random array"""

import os
import random
import logging
import sys
import time
import functools
from multiprocessing import Pool
from itertools import repeat
import numpy as np
from akg.utils.kernel_exec import func_time_required
from akg.utils.kernel_exec import get_profiling_mode

RANDOM_SEED_NUM = 20


def func(size_, miu_=0, sigma_=8, seed_=None):
    """
    Select random func according to RANDOM_FUNC_MODE and randint, calculated by the length of the random_func_list.

    Args:
        size_ (int): Size of data.
        miu_ (int): Mean value. Default: 0.
        sigma_ (int): Standard deviation. Default: 8.
        seed_ (int): seed for random.

    Returns:
        Random func, from random_func_list.
    """
    size_ = (size_ + RANDOM_SEED_NUM - 1) // RANDOM_SEED_NUM
    random_func_list = [
        np.random.RandomState(seed_).normal(miu_, sigma_, size_),
        np.random.RandomState(seed_).logistic(miu_, sigma_, size_),
        np.random.RandomState(seed_).laplace(miu_, sigma_, size_),
        np.random.RandomState(seed_).uniform(miu_, sigma_, size_),
        np.random.RandomState(seed_).tomaxint(size_),
    ]
    env_dic = os.environ
    if not env_dic.get('RANDOM_FUNC_MODE'):
        func_idx = 0
    else:
        func_idx = np.random.RandomState(None).randint(len(random_func_list))
    res = random_func_list[func_idx]
    return res


@func_time_required
def random_gaussian(size, miu=0, sigma=8, epsilon=0, seed=None):
    def random_gaussian_without_bin():
        if sigma <= 0:
            sys.stderr.write("Error: Expect positive sigmal for gaussian distribution. but get %f\n" % sigma)
            return []
        size_c = 1
        for x in size:
            size_c = size_c * x

        if seed is None:
            seed_ = []
            for i in range(RANDOM_SEED_NUM):
                now = int(time.time() % 10000 * 10000) + random.randint(i, 100)
                seed_.append(now)
        else:
            seed_ = [seed] * RANDOM_SEED_NUM
        logging.debug("random_gaussian seeds: {}".format(seed_))
        # In the profiling scenario, when a new process is used to run test cases, data generated by multiple processes
        # stops responding. To locate the fault, please set this parameter gen_data_multi_process to False.
        gen_data_multi_process = not bool(get_profiling_mode())
        if gen_data_multi_process:
            with Pool(processes=8) as pool:
                ret = np.array(pool.starmap(func, zip(repeat(size_c), repeat(miu), repeat(sigma), seed_)))
        else:
            numbers = list()
            for s in seed_:
                numbers.extend(func(size_c, miu, sigma, s))
            ret = np.array(numbers)
        ret = ret.flatten()
        return ret[:size_c].reshape(size) + epsilon

    """Generate random array with absolution value obeys gaussian distribution."""
    random_data_disk_path = None
    if os.environ.get("RANDOM_DATA_DISK_PATH") is not None:
        random_data_disk_path = os.environ.get("RANDOM_DATA_DISK_PATH") + "/random_data_%s_%s.bin" % (
            str(miu), str(sigma))

    if random_data_disk_path is None or (not os.path.exists(random_data_disk_path)):
        return random_gaussian_without_bin()

    data_len = functools.reduce(lambda x, y: x * y, size)
    data_pool = np.fromfile(random_data_disk_path)
    if len(data_pool) == 0:
        return random_gaussian_without_bin()

    if data_len % len(data_pool) != 0:
        copy_num = (data_len // len(data_pool)) + 1
    else:
        copy_num = data_len // len(data_pool)
    data_copy = np.copy(data_pool)
    data_copy_list = []
    for _ in range(copy_num):
        np.random.shuffle(data_copy)
        data_copy_list.append(data_copy)
    data_pool = np.concatenate(tuple(data_copy_list), axis=0)
    return data_pool[0:data_len].reshape(size) + epsilon


def gen_epsilon(dtype):
    """Generate suggested epsilon according to data type."""
    return 1e-7 if dtype == np.float32 else 1e-3


def gen_indices_tensor_scatter_add(shape1, shape2, dtype2):
    """Generate indices for tensor_scatter_add."""
    assert dtype2 == "int32", "Currently only support int32 indices"
    indices = np.zeros(shape2, dtype2)
    indices = indices.reshape(-1, indices.shape[-1])
    for i in range(indices.shape[0]):
        for j in range(indices.shape[1]):
            # add outbounds situation
            if np.random.random() < 0.1:
                # less than 0 and larger than original shape
                indices[i][j] = -1 if np.random.random() < 0.5 else shape1[j] + 10
            else:
                indices[i][j] = np.random.randint(shape1[j], size=())
    indices = indices.reshape(shape2)
    return indices


def gen_indices_gather(shape1, shape2, dtype2, axis):
    """Generate indices for gather."""
    assert dtype2 == "int32", "Currently only support int32 indices"
    indices = np.random.randint(low=0, high=shape1[axis], size=shape2).astype(dtype2)
    return indices


def gen_indices_unsorted_segment_sum(_, shape2, dtype2, num):
    """Generate indices for gather unsorted_segment_sum."""
    # currently only support 1D
    assert dtype2 == "int32", "Currently only support int32 indices"
    return np.random.randint(low=0, high=num, size=shape2).astype(dtype2)


def gen_indices_gather_nd(shape1, shape2, dtype2):
    """Generate indices for gather nd."""
    out_dim1 = 1
    for i in range(len(shape2) - 1):
        out_dim1 = out_dim1 * shape2[i]
    assert dtype2 == "int32", "Currently only support int32 indices"
    indices = np.zeros([shape2[-1], out_dim1]).astype(dtype2)
    for i in range(shape2[-1]):
        # add outbounds situation
        if np.random.random() < 0.1:
            if np.random.random() < 0.5:
                # less than 0
                indices[i] = np.random.randint(low=0, high=shape1[i], size=out_dim1) - 10
            else:
                # larger than original shape
                indices[i] = np.random.randint(low=0, high=shape1[i], size=out_dim1) + 10
        else:
            indices[i] = np.random.randint(low=0, high=shape1[i], size=out_dim1)

    indices = indices.transpose()
    indices = indices.reshape(shape2)
    return indices


def gen_indices(indices_argument):
    """Generate indices."""
    op_name = indices_argument.name
    data_shape = indices_argument.data_shape
    indices_shape = indices_argument.indices_shape
    indices_dtype = indices_argument.indices_dtype
    attrs = indices_argument.attrs
    if op_name == "Gather":
        return gen_indices_gather(data_shape, indices_shape, indices_dtype, attrs)
    elif op_name == "GatherNd":
        return gen_indices_gather_nd(data_shape, indices_shape, indices_dtype)
    elif op_name == "UnsortedSegmentSum":
        return gen_indices_unsorted_segment_sum(data_shape, indices_shape, indices_dtype, attrs)
    assert op_name == "TensorScatterAdd", "Input OP Name Not Known!"
    return gen_indices_tensor_scatter_add(data_shape, indices_shape, indices_dtype)


def gen_csr_indices(indices_argument):
    """Generate csr indices."""
    data_shape = indices_argument.data_shape
    indices_shape = indices_argument.indices_shape
    indices_dtype = indices_argument.indices_dtype
    indptr_choice = np.arange(0, indices_shape[0], dtype=indices_dtype)
    indptr = np.sort(np.random.choice(indptr_choice, data_shape[0] - 1, replace=True))
    indptr = np.concatenate(
        (np.array([0], dtype=indices_dtype), indptr, np.array([indices_shape[0]], dtype=indices_dtype)))
    indices_choice = np.arange(data_shape[1], dtype=indices_dtype)
    idx_range = indices_choice.copy()
    idx_bound = np.diff(indptr).reshape(-1, 1)
    max_count = idx_bound.max().tolist()
    np.random.shuffle(indices_choice)
    indices_choice = indices_choice[: max_count]
    indices_choice.sort()
    indices = np.where(idx_range[: max_count] < idx_bound, indices_choice, data_shape[1])
    mask = np.less(indices, data_shape[1]).nonzero()
    indices = indices[mask]
    return indptr, indices
